import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import scipy.io as sio
import scipy.sparse as sp
import scipy.sparse.linalg as lg
from time import time

def loadGraphFromEdgeListTxt(file_name, directed=True):
  with open(file_name, 'r') as f:
    # n_nodes = f.readline()
    # f.readline() # Discard the number of edges
    if directed:
      G = nx.DiGraph()
    else:
      G = nx.Graph()
    for line in f:
      edge = line.strip().split()
      if len(edge) == 3:
        w = float(edge[2])
      else:
        w = 1.0
      G.add_edge(edge[0], edge[1], weight=w)
  return G

class HOPE:  
  def __init__(self,d,beta):
    self._d = d
    self._beta = beta



  def learn_embedding(self, graph=None, edge_f=None,
                      is_weighted=False, no_python=False):
    if not graph and not edge_f:
      raise Exception('graph/edge_f needed')
    if not graph:
      graph = loadGraphFromEdgeListTxt(edge_f)

    t1 = time()
    # A = nx.to_scipy_sparse_matrix(graph)
    # I = sp.eye(graph.number_of_nodes())
    # M_g = I - self._beta*A
    # M_l = self._beta*A
    A = nx.to_numpy_matrix(graph)
    # M_g = np.eye(graph.number_of_nodes()) - self._beta * A
    # M_l = self._beta * A
    # S = np.dot(np.linalg.inv(M_g), M_l)
    S = A
    # print S.shape
    # print S[:,1]
    # U,sigma,VT = np.linalg.svd(S)
    # u = U[:,:2]
    # vt = VT[:2,:]
    # s = sigma[:2]
    u, s, vt = lg.svds(S, k=self._d // 2)
    X1 = np.dot(u, np.diag(np.sqrt(s)))
    X2 = np.dot(vt.T, np.diag(np.sqrt(s)))
    t2 = time()
    self._X = np.concatenate((X1, X2), axis=1)

    # p_d_p_t = np.dot(u, np.dot(np.diag(s), vt))
    # eig_err = np.linalg.norm(p_d_p_t - S)
    # print('SVD error (low rank): %f' % eig_err)
    return self._X

def main():
  dG = loadGraphFromEdgeListTxt("../test.g", directed=False)
  dG = dG.to_directed()
  print(dG.nodes())
  hope = HOPE(4,0.01)
  hope.learn_embedding(graph=dG,is_weighted=True)
  node_pos = hope._X[:, :2]
  node_pos = np.squeeze(np.asarray(node_pos))
  # plt.scatter(node_pos[:,0], node_pos[:,1])
  node_num, embedding_dimension = node_pos.shape
  pos = {}
  i = 0
  for n in dG.nodes():
    pos[n] = node_pos[i, :]
    i = i + 1
  print(pos)
  np.save('node_pos.dat',pos) 
  nx.draw_networkx(dG, pos, labels=None)
    # , node_color=None,
    #                          width=0.1, node_size=300, arrows=False,
    #                          alpha=0.8, font_size=12, labels=None)
  plt.show()

if __name__ == '__main__':
  main()  